#ifndef __STSKIPLISTSET_H__
#define __STSKIPLISTSET_H__

#include "BitTwister.h"
//#include <string.h>
#include <sstream>

namespace st
{
	// Forward declaration of the iterator
	template<typename Tkey> class STSkipListSetIterator;

	typedef int levelType;

	template <typename Tkey>
	class STSkipListSet
	{
	protected: 
		struct node
		{
		protected:
			Tkey		_key;
			levelType	_level;
			node**		_next;	// ptrs to next nodes

		public:
			node(const Tkey& key, levelType level):
				_key(key), _level(level)
			{
				_next = new node*[_level]; // 'new' doesn't zero out
				memset(_next, 0, _level * sizeof(node*));
			}

			virtual ~node() { delete[] _next; }
			
			levelType		Level() const				{ return _level; }
			const Tkey&		Key() const					{ return _key; }
			void			SetKey(const Tkey& key)		{ _key = key; }
			node*			Next(levelType level) const	{ return _next[level]; }
			void			SetNext(levelType level, node* aNode) { _next[level] = aNode; }
			void			Clear()						{ memset(_next, 0, _level * sizeof(node*)); }
		};

		node*		_head;
		levelType	_maxLevel; // see comment in the constructor
		levelType	_level;
		BitTwister*	_rand;
		size_t		_count;

		levelType RandLevel();
		bool ValidLevel(levelType level) const;
		void RemoveAll();

	public:
		// Befriend the iterator so it can access the nodes directly
		friend class STSkipListSetIterator<Tkey>;
		typedef STSkipListSetIterator<Tkey> iterator;

		// Level should be the binary power of the maximum number of elements
		// expected. So for ~30K elements, 2^16 = 32K, max level should be 16.
		STSkipListSet(levelType maxLevel = 16):
			_maxLevel(maxLevel + 1), _level(0), _count(0)
		{
			_head = new node(0, _maxLevel);
			_rand = new BitTwister();
		}

		virtual ~STSkipListSet() { RemoveAll(); delete _rand; delete _head; }

		void Insert(const Tkey& key);
		iterator Search(const Tkey& key); // returns 'end()' if not found
		iterator SearchLessEqual(const Tkey& key) const; // find the last node that has a valeu less or equal than 'key', end() if not found
		void Remove(const Tkey& key);

		iterator	begin()	{ return iterator(_head->Next(0), this); }
		iterator	end() const		{ return NULL; }
		size_t		size() const	{ return _count; }
		void		clear()			{ RemoveAll(); }

		bool Check() const;
		virtual void Stats(std::stringstream& sout) const;
		void StatsStdout() const; // same as above, it'll just dump it to stdout
	};

	template<typename Tkey>
	class STSkipListSetIterator: public std::iterator<std::forward_iterator_tag, int>
	{
		typedef typename STSkipListSet<Tkey>::node STSkipListSetNode;
		typedef STSkipListSetIterator<Tkey> thisType;

		STSkipListSetNode*			_currentNode;
		STSkipListSet<Tkey>*		_parentList;

	public:
		STSkipListSetIterator() : _currentNode(NULL), _parentList(NULL) {}
		STSkipListSetIterator(STSkipListSetNode* node): _currentNode(node), _parentList(NULL) {}
		STSkipListSetIterator(STSkipListSetNode* node, STSkipListSet<Tkey>* parent): _currentNode(node), _parentList(parent) {}
		STSkipListSetIterator(const thisType& iter): _currentNode(iter._currentNode), _parentList(iter._parentList) {}
 
		bool operator==(const thisType& right) const
		{
			return _currentNode == right._currentNode;
		}
 
		bool operator!=(const thisType& right) const
		{
			return _currentNode != right._currentNode;
		}
 
		thisType& operator++()
		{
			if(_currentNode != NULL)
				_currentNode = _currentNode->Next(0);
	        return *this;
		}
 
		void iterateToNext()
		{
			if(_currentNode != NULL)
				_currentNode = _currentNode->Next(0);
		}

		const Tkey& operator*() const { return _currentNode->Key(); }
		Tkey* operator->() const { return &(_currentNode->Key()); }
 
		operator void*() const { return _currentNode; }

		// These should be first/second to mimic 'map' but I'm not using std::pair
		const Tkey& Key() const { STASSERT(_currentNode); return _currentNode->Key(); }
	};

	template <typename Tkey>
	levelType STSkipListSet<Tkey>::RandLevel()
	{
		levelType level = 1;
		levelType rand = _rand->bitRandomInt(_maxLevel); // only 6 random bits needed for a max of 64
		// This is equivalent to p = 1/2 in Pugh's paper but it won't work if _maxLevel > 32
		while((rand & 0x1) && level < _maxLevel)
		{
			level++;
			rand >>= 1;
		}
		// Fix the dice proposition from Pugh's paper
		if(level > _level + 1)
			level = _level + 1;
		
		return level;
	}

	template <typename Tkey>
	STSkipListSetIterator<Tkey> STSkipListSet<Tkey>::Search(const Tkey& key)
	{
		if(_head == NULL)
			return end();

		node* head = _head;

		for(levelType level = _level - 1; level >= 0; level--)
			while(head->Next(level) != NULL && head->Next(level)->Key() < key)
				head = head->Next(level);

		head = head->Next(0);

		if(head != NULL && head->Key() == key)
			return STSkipListSetIterator<Tkey>(head, this);
		else
			return end();
	}

	template <typename Tkey>
	STSkipListSetIterator<Tkey> STSkipListSet<Tkey>::SearchLessEqual(const Tkey& key) const
	{
		if(_head == NULL)
			return end();

		node* head = _head;

		for(levelType level = _level - 1; level >= 0; level--)
			while(head->Next(level) != NULL && head->Next(level)->Key() < key)
				head = head->Next(level);

		if(head->Next(0) != NULL && head->Next(0)->Key() <= key)
			head = head->Next(0);

		if(head != NULL && head != _head)
			return iterator(head, this);
		else
			return end();
	}

	template <typename Tkey>
	void STSkipListSet<Tkey>::Remove(const Tkey& key)
	{
		STASSERT(_head && _level <= _maxLevel);

		if(_head == NULL)
			return;

		node* head = _head;
		node** update = new node*[_maxLevel];
		memset(update, 0, sizeof(node*) * _maxLevel);
		
		for(levelType level = _level - 1; level >= 0; level--)
		{
			while(head->Next(level) != NULL && head->Next(level)->Key() < key)
				head = head->Next(level);
			update[level] = head;
		}

		head = head->Next(0);

		if(head != NULL && head->Key() == key)
		{
			for(levelType level = 0; level < head->Level(); level++)
			{
				if(update[level]->Next(level) == head)
					update[level]->SetNext(level, head->Next(level));
			}
			delete head;
			while(_level > 0 && _head->Next(_level - 1) == NULL)
				_level--;
			_count--;
		}
		delete[] update;
	}

	template <typename Tkey>
	void STSkipListSet<Tkey>::Insert(const Tkey& key)
	{
		node* head = _head;
		node** update = new node*[_maxLevel];
		memset(update, 0, sizeof(node*) * _maxLevel);

		for(levelType level = _level - 1; level >= 0; level--)
		{
			while(head->Next(level) != NULL && head->Next(level)->Key() < key)
				head = head->Next(level);
			update[level] = head;
		}

		head = head->Next(0);

		if(head != NULL && head->Key() == key)
		{
			// do nothing, duplicate key
		}
		else
		{
			levelType newLevel = RandLevel();
			if(newLevel > _level)
			{
				for(levelType level = _level; level < newLevel; level++)
					update[level] = _head;
				_level = newLevel;
			}
			node* newNode = new node(key, newLevel);
			for(levelType level = 0; level < newLevel; level++)
			{
				newNode->SetNext(level, update[level]->Next(level));
				update[level]->SetNext(level, newNode);
			}
			_count++;
		}

		delete[] update;
	}

	template <typename Tkey>
	bool STSkipListSet<Tkey>::ValidLevel(levelType level) const
	{
		// Levels above 31 don't work due to 4 byte ints (2^32)
		return _maxLevel <= 30 && level >= 0 && level <= _maxLevel;
	}

	template <typename Tkey>
	void STSkipListSet<Tkey>::RemoveAll()
	{
		if(!_head)
			return;
		node* head = _head->Next(0);
		node* next;
		while(head != NULL)
		{
			next = head->Next(0);
			delete head;
			head = next;
		}
		_head->Clear();
		_level = 0;
		_count = 0;
	}

	template <typename Tkey>
	bool STSkipListSet<Tkey>::Check() const
	{
		node* head = _head->Next(0); // skip the head
		size_t count = 0;
		while(head != NULL)
		{
			count++;
			if(head->Level() > _level || head->Level() < 1)
				return false;
			for(levelType l = 0; l < head->Level(); l++)
				if(head->Next(l) != NULL && head->Next(l)->Key() <= head->Key())
					return false;
			head = head->Next(0);
		}
		if(count != _count)
			return false;
		else
			return true;
	}

	template <typename Tkey>
	void STSkipListSet<Tkey>::Stats(std::stringstream& sout) const
	{
		sout << "Total elements in list: " << _count << std::endl;
		if(_level > 0 && _count > 0)
		{
			levelType* levelsCount = new levelType[_level];
			memset(levelsCount, 0, sizeof(levelType) * _level);
			
			node* head = _head->Next(0);
			while(head != NULL)
			{
				levelsCount[head->Level() - 1]++;
				head = head->Next(0);
			}
			for(levelType i = 0; i < _level; i++)
				sout << "Level " << i + 1 << ": " << levelsCount[i] << std::endl;
			delete[] levelsCount;
		}
	}

	template <typename Tkey>
	void STSkipListSet<Tkey>::StatsStdout() const
	{
		std::stringstream sstats;
		Stats(sstats);
		std::cout << sstats.str();
	}
}

#endif // __STSKIPLISTSET_H__
